#!/usr/bin/env python3
# Copyright (c) 2019-2020 Windy Albert

import hashlib
import base64
import binascii
import hmac
import ecdsa
import os
import base58check
from ecdsa.ellipticcurve import AbstractPoint, PointJacobi
import WordList
import random

COLOR_NONE  = "\033[0m"
COLOR_RED   = "\033[0;31m"
COLOR_GREEN = "\033[0;32m"
COLOR_YELLOW = "\033[0;33m"


class keyTree(object):
    version_xprv = "0488ade4" # xprv
    version_xpub = "0488b21e" # xpub
    version_yprv = "049d7878" # yprv (for extended keys in a BIP 49 derivation path)
    version_ypub = "049d7cb2" # ypub
    version_zprv = "04b2430c" # zprv (for extended keys in a BIP 84 derivation path)
    version_zpub = "04b24746" # zpub

    def __init__(self, mnemonic=""):
        self.mnemonic = mnemonic

    def getMasterKey(self,password=""):
        print("mnemonic               : "+ COLOR_RED + self.mnemonic +COLOR_NONE)
        seed = hashlib.pbkdf2_hmac('sha512', self.mnemonic.encode('utf-8'), b'mnemonic'+password, 2048)
        print("seed                   : "+COLOR_RED+str.format(seed.hex())+COLOR_NONE)
        hash0 = hmac.new(b'Bitcoin seed', bytes.fromhex(seed.hex()),  digestmod='sha512')        
        prikey    = hash0.hexdigest()[0:64]
        chaincode = hash0.hexdigest()[64:128]
        fingerprint = "00000000"
        index = "00000000"
        depth = "00"
        xprv = self.version_xprv+depth+fingerprint+index+chaincode+"00"+prikey
        sha256_1 = hashlib.new('sha256',bytes.fromhex(xprv))
        sha256_2 = hashlib.new('sha256',bytes.fromhex(sha256_1.hexdigest()))
        xprv_b58 = base58check.b58encode(bytes.fromhex(xprv+(sha256_2.hexdigest())[0:8]))

        result = {}
        result["path"] = "m"
        result["xprv"] = xprv_b58.decode('utf-8')
        print(result["path"] + " "*(23-len(result["path"]))+": " + COLOR_RED + result["xprv"] + COLOR_NONE)

        return result

    def getChildpriKey(self,obj,pa):
        result = {}
        dep = len(obj["path"].split("/"))
        xprv_hex = base58check.b58decode( obj["xprv"] )
        chaincode = xprv_hex[13:45]
        prikey    = xprv_hex[46:78]

        sk = ecdsa.SigningKey.from_string( bytes.fromhex(prikey.hex()), curve=ecdsa.SECP256k1)
        vk = sk.get_verifying_key() #this is your verification key (public key)
        public_key = vk.to_string("compressed").hex()
        fingerprint = self.__getFingerprint(public_key)

        if("\'" in pa):
            index = '{:08x}'.format(int(pa[0:-1])+0x80000000)
            private_key = sk.to_string().hex() #convert your private key to hex
            hmacsha512 = hmac.new(chaincode, bytes.fromhex('00'+private_key+index),  digestmod='sha512')
        else:
            index = '{:08x}'.format(int(pa))
            hmacsha512 = hmac.new(chaincode, bytes.fromhex(public_key      +index),  digestmod='sha512')

        left  = hmacsha512.hexdigest()[ 0:64]
        right = hmacsha512.hexdigest()[64:128]
        
        h = int(prikey.hex(),16)
        i = int(left  ,16)
        
        childkey = (h+i)%ecdsa.SECP256k1.generator.order()
        st = childkey.to_bytes(32,byteorder='big',signed=False)

        prikey_child = st.hex()
        chaincode_child = right

        depth = '{:02x}'.format(int(dep))
        
        xprv_c = self.version_xprv+depth+fingerprint+index+chaincode_child+"00"+prikey_child
        sha256_1 = hashlib.new('sha256',bytes.fromhex(xprv_c))
        sha256_2 = hashlib.new('sha256',bytes.fromhex(sha256_1.hexdigest()))
        xprv = base58check.b58encode(bytes.fromhex(xprv_c+(sha256_2.hexdigest())[0:8]))
        
        if dep > 2:
            sk_child = ecdsa.SigningKey.from_string( bytes.fromhex(st.hex()), curve=ecdsa.SECP256k1)
            vk_child = sk_child.get_verifying_key() #this is your verification key (public key)
            pubkey_child = vk_child.to_string("compressed").hex()

            xpub_c = self.version_xpub+depth+fingerprint+index+chaincode_child+pubkey_child
            sha256_1 = hashlib.new('sha256',bytes.fromhex(xpub_c))
            sha256_2 = hashlib.new('sha256',bytes.fromhex(sha256_1.hexdigest()))
            xpub = base58check.b58encode(bytes.fromhex(xpub_c+(sha256_2.hexdigest())[0:8]))
            result["xpub"] = xpub.decode('utf-8')
        else :
            result["xpub"] = ""
        
        result["path"] = obj["path"] + "/" + pa
        result["xprv"] = xprv.decode('utf-8')
        if (dep < 5):
            print( result["path"] + " "*(23-len(result["path"])) + ": "+ COLOR_RED + result["xprv"] + COLOR_NONE)
        return result



    def getChildpubKey(self,obj,pa):
        result = {}
        dep = len(obj["path"].split("/"))
        xpub_hex = base58check.b58decode( obj["xpub"] )
        
        sha256_1 = hashlib.new('sha256',xpub_hex[0:-4])
        sha256_2 = hashlib.new('sha256',bytes.fromhex(sha256_1.hexdigest()))
        cs1 = sha256_2.hexdigest()[0:8]
        cs2 = xpub_hex[-4:].hex()
        if cs1 != cs2 :
            print("Pubkey is broken ! Cancel ...")
            return
            
        chaincode = xpub_hex[13:45]
        public_key= xpub_hex[45:78].hex()

        fingerprint = self.__getFingerprint( public_key )

        index = '{:08x}'.format(int(pa))
        hmacsha512 = hmac.new(chaincode, bytes.fromhex(public_key+index),  digestmod='sha512')

        left  = hmacsha512.hexdigest()[ 0: 64]
        right = hmacsha512.hexdigest()[64:128]
        
        sk_1 = ecdsa.SigningKey.from_string( bytes.fromhex(left), curve=ecdsa.SECP256k1)        
        vk_1 = sk_1.get_verifying_key().to_string("uncompressed").hex()        
        x1 = int(vk_1[2:66],16)
        y1 = int(vk_1[66:130],16)
        point1 = PointJacobi(curve=ecdsa.SECP256k1.curve, x=x1, y=y1, z=1, order=ecdsa.SECP256k1.generator.order(), generator=False)
        
        vk_2 = ecdsa.VerifyingKey.from_string( bytes.fromhex(public_key), curve=ecdsa.SECP256k1)
        vk_2 = vk_2.to_string("uncompressed").hex()     
        x2 = int(vk_2[2:66],16)
        y2 = int(vk_2[66:130],16)
        point2 = PointJacobi(curve=ecdsa.SECP256k1.curve, x=x2, y=y2, z=1, order=ecdsa.SECP256k1.generator.order(), generator=False)
        
        point3 = point1.__add__(point2)

        x3 = '{:064x}'.format(point3.x())
        y3 = point3.y()
        pubkey_child = '03'+x3  if (y3%2)==1 else '02'+x3 # compressed


        chaincode_child = right
        depth = '{:02x}'.format(int(dep))

        xpub_c = self.version_xpub+depth+fingerprint+index+chaincode_child+pubkey_child
        sha256_1 = hashlib.new('sha256',bytes.fromhex(xpub_c))
        sha256_2 = hashlib.new('sha256',bytes.fromhex(sha256_1.hexdigest()))
        xpub = base58check.b58encode(bytes.fromhex(xpub_c+(sha256_2.hexdigest())[0:8]))
        result["path"] = obj["path"] + "/" + pa
        result["xpub"] = xpub.decode('utf-8')
        
        if (dep < 5):
            print( result["path"] + " "*(23-len(result["path"])) + ": "+ COLOR_GREEN + result["xpub"] + COLOR_NONE)
        return result,pubkey_child


    def __getFingerprint(self,key):
        hash_256 = hashlib.sha256()
        hash_256.update(bytes.fromhex(key))
        hash_256_value = hash_256.hexdigest()
        obj = hashlib.new('ripemd160', bytes.fromhex(hash_256_value))
        ripemd160 = obj.hexdigest()
        fingerprint = ripemd160[0:8]
        return fingerprint

def getRandomMemonic(rand):
    s = '{:040x}'.format(rand)

    sha256  = hashlib.new('sha256',bytes.fromhex(s))
    checksum = sha256.hexdigest()
    b = bytes.fromhex(checksum)
    index = (rand<<5) | (b[0]>>(8-5))

    data =""
    for i in range(15):
        if(i != 0):
            data = data + " "
        n = (index>>(165 - 11 - i * 11)) & 0x07FF
        data = data +WordList.WORDLISTS_en[n]
        
    return {"rand":s,"mnemonic":data}

def print_line(text,color):
    print(color + "═"*136 + COLOR_NONE)
    print(color + " "*62 + " " + text + " " + " "*(72-len(text)) + COLOR_NONE)




def entropy_from_mnemonic(mnemonic):
    words = mnemonic.split(" ")
    tbl = {}
    for i in range(2048):
        data = WordList.WORDLISTS_en[i]
        tbl[data] = i

    entropy = 0
    for wo in words:
        n = tbl[wo]
        entropy = (entropy << 11) | n
    entropy = entropy >> 4
    
    return '{:32x}'.format(int(entropy))

def main():
    # mnemonic = getRandomMemonic()
    mnemonic = 'forest bag bleak group million equal problem pen surge crowd approve name'

    tree = keyTree(mnemonic)
    password = b''
    xprv_m = tree.getMasterKey(password)    
    xprv_m_44h = tree.getChildpriKey(xprv_m,"44\'")
           

if __name__ == '__main__':
    main()